package simple.demo.zk;

import java.io.IOException;
import java.util.concurrent.CountDownLatch;

import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.ZKUtil;
import org.apache.zookeeper.ZooKeeper;
import org.apache.zookeeper.ZooDefs.Ids;

public class ZkLock {
	
	private final String name;
	private ZooKeeper zk;
	private final Object mutex = new Object();
	private CountDownLatch connectLatch = new CountDownLatch(1);
	
	private final Watcher watcher = new Watcher() {
		
		@Override
		public void process(org.apache.zookeeper.WatchedEvent event) {
			if(event.getState() == Watcher.Event.KeeperState.SyncConnected)
				connectLatch.countDown();
			
			synchronized (mutex) {
				mutex.notify();
			}
		};
	};
	
	public ZkLock(String name) {
		this.name = name;
	}
	
	public void connectTo(String host) throws IOException {
		zk = new ZooKeeper(host, 3000, watcher);
	}
	
	public void createRoot() throws KeeperException, InterruptedException {
		zk.create(ZkLockHelper.getLockRoot(), null, Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
	}
	
	void cleanUp() throws InterruptedException, KeeperException {
		ZKUtil.deleteRecursive(zk, ZkLockHelper.getLockRoot());
	}
	
	public void waitForConnection() throws InterruptedException {
		connectLatch.await();
	}
	
	public boolean lock() {
		while(true) {
			try {
				if(zk.exists(ZkLockHelper.getLockPath(name), true) != null)
				{
					synchronized (mutex) {
						mutex.wait();
					}
				}
			} catch (InterruptedException ex) {
				return false;
			} catch (KeeperException e) {
				return false;
			} 
			
			try {
				zk.create(ZkLockHelper.getLockPath(name), null, Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL);
				return true;
			} catch(KeeperException ex) {
				//continue
			} catch (InterruptedException e) {
				return false;
			}
		}
	}
	
	public void unlock() throws InterruptedException, KeeperException {
		zk.delete(ZkLockHelper.getLockPath(name), 0);
	}

}
